# coding=utf-8
import matplotlib.pyplot as plt
import pandas as pd
import tushare as ts


# 计算KDJ
def calKDJ(df):
    df['MinLow'] = df['low'].rolling(9, min_periods=9).min()
    # 填充NaN数据
    df['MinLow'].fillna(value=df['low'].expanding().min(), inplace=True)
    df['MaxHigh'] = df['high'].rolling(9, min_periods=9).max()
    df['MaxHigh'].fillna(value=df['high'].expanding().max(), inplace=True)
    df['RSV'] = (df['close'] - df['MinLow']) / (df['MaxHigh'] - df['MinLow']) * 100
    # 通过for循环依次计算每个交易日的KDJ值
    for i in range(len(df)):
        if i == 0:  # 第一天
            df.loc[i, 'K'] = 50
            df.loc[i, 'D'] = 50
        if i > 0:
            df.loc[i, 'K'] = df.loc[i - 1, 'K'] * 2 / 3 + 1 / 3 * df.loc[i, 'RSV']
            df.loc[i, 'D'] = df.loc[i - 1, 'D'] * 2 / 3 + 1 / 3 * df.loc[i, 'K']
        df.loc[i, 'J'] = 3 * df.loc[i, 'K'] - 2 * df.loc[i, 'D']
    return df


def drawKDJ():
    pro = ts.pro_api('183e551b07ab171b5570a8899bb3c8b263b517981af9d1d8faf86f6a')

    # 获取指定股票开始日期和结束日期之间的日线行情
    df = pro.daily(ts_code='000001.SZ', start_date='20180901', end_date='20190531')
    # df = pd.read_csv('D:/stockData/ch8/6035052018-09-012019-05-31.csv', encoding='gbk')
    stockDataFrame = calKDJ(df)
    print(stockDataFrame)
    # 开始绘图
    plt.figure()
    stockDataFrame['K'].plot(color="blue", label='K')
    stockDataFrame['D'].plot(color="green", label='D')
    stockDataFrame['J'].plot(color="purple", label='J')
    plt.legend(loc='best')  # 绘制图例
    # 设置x轴坐标的标签和旋转角度
    major_index = stockDataFrame.index[stockDataFrame.index % 10 == 0]
    major_xtics = stockDataFrame['trade_date'][stockDataFrame.index % 10 == 0]
    plt.xticks(major_index, major_xtics)
    plt.setp(plt.gca().get_xticklabels(), rotation=30)
    # 带网格线，且设置了网格样式
    plt.grid(linestyle='-.')
    plt.title("000001.SZ的KDJ图")
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.show()


# 调用方法
drawKDJ()
